import open3d as o3d
import numpy as np
from map_trans_toolkit_pkg.scripts.feature_extract_pcd import extract_features

# ==== 参数配置 ====
pcd_file1 = "../maps/floor7_clip.pcd"
pcd_file2 = "../maps/stair1_clip.pcd"
output_file = "../maps/merged_pcd.pcd"
prior_tf = [6.75, -7.5, -0.17, 0, 0, -1.63]
voxel_size = 0.05
random_sample_ratio = 0.2
# 特征提取参数
feature_voxel = 0.05
feature_normal_radius = 0.15
feature_normal_max_nn = 40
feature_curvature_knn = 15
feature_angle_threshold = 0.4
feature_curvature_threshold = 0.3
# ==========================


def euler_to_matrix(x, y, z, roll, pitch, yaw):
    Rx = np.array(
        [[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]]
    )
    Ry = np.array(
        [
            [np.cos(pitch), 0, np.sin(pitch)],
            [0, 1, 0],
            [-np.sin(pitch), 0, np.cos(pitch)],
        ]
    )
    Rz = np.array(
        [[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]]
    )
    R = Rz @ Ry @ Rx
    T = np.eye(4)
    T[:3, :3] = R
    T[:3, 3] = [x, y, z]
    return T


def crop_by_aabb(pcd, min_bound, max_bound):
    aabb = o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound)
    return pcd.crop(aabb)


if __name__ == "__main__":
    print("[1/8] 读取点云...")
    pcd1 = o3d.io.read_point_cloud(pcd_file1)
    pcd2 = o3d.io.read_point_cloud(pcd_file2)

    print("[2/8] 提取特征点云...")
    feat1 = extract_features(
        pcd1,
        voxel_size=feature_voxel,
        normal_radius=feature_normal_radius,
        normal_max_nn=feature_normal_max_nn,
        curvature_knn=feature_curvature_knn,
        angle_threshold=feature_angle_threshold,
        curvature_threshold=feature_curvature_threshold,
    )
    feat2 = extract_features(
        pcd2,
        voxel_size=feature_voxel,
        normal_radius=feature_normal_radius,
        normal_max_nn=feature_normal_max_nn,
        curvature_knn=feature_curvature_knn,
        angle_threshold=feature_angle_threshold,
        curvature_threshold=feature_curvature_threshold,
    )
    print(
        f"  特征点数1: {np.asarray(feat1.points).shape[0]}, 特征点数2: {np.asarray(feat2.points).shape[0]}"
    )

    print("[3/8] 对第二个点云及其特征点应用先验变换...")
    tf = euler_to_matrix(*prior_tf)
    pcd2.transform(tf)
    feat2.transform(tf)

    print("[4/8] 基于特征点云进行粗配准...")
    icp_init = o3d.pipelines.registration.registration_icp(
        feat2,
        feat1,
        2.0,
        np.eye(4),
        o3d.pipelines.registration.TransformationEstimationPointToPoint(),
    )
    print(f"  粗配准 transformation:\n{icp_init.transformation}")

    pcd2.transform(icp_init.transformation)
    feat2.transform(icp_init.transformation)

    print("[5/8] 对原始点云做随机降采样，进行精配准...")
    pcd1_down = pcd1.random_down_sample(random_sample_ratio)
    pcd2_down = pcd2.random_down_sample(random_sample_ratio)
    icp_fine = o3d.pipelines.registration.registration_icp(
        pcd2_down,
        pcd1_down,
        1.0,
        np.eye(4),
        o3d.pipelines.registration.TransformationEstimationPointToPoint(),
    )
    print(f"  精配准 transformation:\n{icp_fine.transformation}")

    pcd2.transform(icp_fine.transformation)

    print("[6/8] 合并点云...")
    merged = pcd1 + pcd2

    print("[7/8] 体素滤波降采样...")
    try:
        merged_down = merged.voxel_down_sample(voxel_size)
        print(f"  降采样后点数: {np.asarray(merged_down.points).shape[0]}")
    except RuntimeError as e:
        print(
            f"  [警告] voxel_size={voxel_size} 太小，降采样失败。自动尝试更大的体素尺寸..."
        )
        bbox = merged.get_axis_aligned_bounding_box()
        diag = np.linalg.norm(
            np.array(bbox.get_max_bound()) - np.array(bbox.get_min_bound())
        )
        auto_voxel = max(0.2, diag / 1000)
        print(f"  自动设置voxel_size={auto_voxel:.3f}，重新降采样...")
        merged_down = merged.voxel_down_sample(auto_voxel)
        print(f"  降采样后点数: {np.asarray(merged_down.points).shape[0]}")

    print("[8/8] 保存并可视化...")
    o3d.io.write_point_cloud(output_file, merged_down)
    print(f"合并点云已保存到: {output_file}")
    o3d.visualization.draw_geometries([merged_down])
